// Package saramax
// Created by GoLand
// @User: lenora
// @Date: 2024/7/12
// @Time: 17:48

package saramax

import (
	"camp/microbook/pkg/logger"
	"context"
	"encoding/json"
	"github.com/IBM/sarama"
	"golang.org/x/sync/errgroup"
	"time"
)

type Handler[T any] struct {
	fn func(message *sarama.ConsumerMessage, data T) error
	l  logger.Logger
}

func NewHandler[T any](fn func(message *sarama.ConsumerMessage, data T) error, l logger.Logger) *Handler[T] {
	return &Handler[T]{fn: fn, l: l}
}

type HandlerMsg[T any] struct {
	message *sarama.ConsumerMessage
	data    T
}

func (h *Handler[T]) Setup(session sarama.ConsumerGroupSession) error {
	return nil
}

func (h *Handler[T]) Cleanup(session sarama.ConsumerGroupSession) error {
	return nil
}

func (h *Handler[T]) ConsumeClaim(session sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
	msgs := claim.Messages()
	const batchSize = 10
	for {
		batch := make([]*sarama.ConsumerMessage, 0, batchSize)
		ts := make([]T, 0, batchSize)
		ctx, cancel := context.WithTimeout(context.TODO(), time.Second*10)
		done := false
		for i := 0; i < batchSize && !done; i++ {
			select {
			case <-ctx.Done():
				done = true
			case msg, ok := <-msgs:
				if !ok {
					cancel()
					return nil
				}
				h.l.Info("sarama msg",
					logger.String("topic", msg.Topic),
					logger.Int32("partition", msg.Partition),
					logger.Int64("offset", msg.Offset),
				)
				var data T
				err := json.Unmarshal(msg.Value, &data)
				if err != nil {
					h.l.Error("decode msg failed",
						logger.Error(err),
						logger.String("topic", msg.Topic),
						logger.Int32("partition", msg.Partition),
						logger.Int64("offset", msg.Offset),
					)
					continue
				}
				batch = append(batch, msg)
				ts = append(ts, data)
			}
		}
		cancel()
		var eg errgroup.Group
		for i, message := range batch {
			eg.Go(func() error {
				err := h.fn(message, ts[i])
				if err != nil {
					h.l.Error("handle msg failed",
						logger.Error(err),
						logger.String("topic", message.Topic),
						logger.Int32("partition", message.Partition),
						logger.Int64("offset", message.Offset),
					)
				}
				return nil
			})
		}
		_ = eg.Wait()
		for _, msg := range batch {
			session.MarkMessage(msg, "")
		}
	}
}
